import numpy as np
import cv2 as cv


# 导入numpy库用于数值计算，导入OpenCV库并简称为cv，用于后续的图像处理操作


# 灰度图转换
def grayscale(image):
    return cv.cvtColor(image, cv.COLOR_RGB2GRAY)


# 定义函数grayscale，功能是将输入的彩色图像（RGB格式）转换为灰度图像，
# 通过调用OpenCV的cvtColor函数，并指定转换的颜色空间模式为从RGB到GRAY


# 高斯滤波
def gaussian_blur(image, kernel_size):
    return cv.GaussianBlur(image, (kernel_size, kernel_size), 0)


# 定义函数gaussian_blur，用于对输入图像进行高斯滤波操作，
# 接收图像和卷积核大小（kernel_size）作为参数，
# 调用OpenCV的GaussianBlur函数，其中最后一个参数0表示在x和y方向上的标准差自动计算


# Canny边缘检测
def canny(image, low_threshold, high_threshold):
    return cv.Canny(image, low_threshold, high_threshold)


# 定义函数canny，用于对输入图像进行Canny边缘检测，
# 接收图像以及低阈值和高阈值作为参数，调用OpenCV的Canny函数实现边缘检测功能


# 生成感兴趣区域即Mask掩模,从输入图像中提取出我们感兴趣的特定区域，同时将图像的其余部分屏蔽掉
def region_of_interest(image, vertices):
    mask = np.zeros_like(image)  # 生成图像大小一致的zeros矩
    # 创建一个与输入图像大小相同、元素全为0的数组，用于后续构建感兴趣区域的掩模


    # 填充顶点vertices中间区域
    if len(image.shape) > 2:
        channel_count = image.shape[2]
        ignore_mask_color = (255,) * channel_count
        # 如果图像是彩色的（有多个通道），则创建一个与通道数对应的元组，元素都为255，用于填充掩模
    else:
        ignore_mask_color = 255
        # 如果图像是灰度图（单通道），则直接将填充颜色设为255


    # 填充函数
    cv.fillPoly(mask, vertices, ignore_mask_color)
    # 使用OpenCV的fillPoly函数，根据给定的顶点（vertices）信息，用指定颜色（ignore_mask_color）填充多边形区域，构建掩模
    masked_image = cv.bitwise_and(image, mask)
    # 按位与，只有当对应位置的 image 和 mask 像素值都不为 0 时，结果才不为 0，也就是该像素保留显示；若 mask 中对应像素为 0，则结果为 0，对应原图像中的像素就被屏蔽掉不显示了
    return masked_image


# 原图像与车道线图像按照a:b比例融合
def weighted_img(img, initial_img, a=0.8, b=1., c=0.):
    return cv.addWeighted(initial_img, a, img, b, c)


# 定义函数weighted_img，用于将两张图像按照给定的权重系数进行融合，
# 调用OpenCV的addWeighted函数，其中a、b、c分别为对应图像的权重系数及偏置项


def hough_lines(img, rho, theta, threshold, min_line_len, max_line_gap):
    # rho：线段以像素为单位的距离精度
    # theta : 像素以弧度为单位的角度精度(np.pi/180较为合适)
    # threshold : 霍夫平面累加的阈值
    # minLineLength : 线段最小长度(像素级)
    # maxLineGap : 最大允许断裂长度
    lines = cv.HoughLinesP(img, rho, theta, threshold, np.array([]), minLineLength=min_line_len,
                       maxLineGap=max_line_gap)
    return lines


# 定义函数hough_lines，用于在输入图像上应用霍夫变换来检测直线，
# 接收霍夫变换相关的参数（rho、theta、threshold、min_line_len、max_line_gap），
# 调用OpenCV的HoughLinesP函数进行直线检测，并返回检测到的直线信息


def draw_lines(image, lines, color=[255, 0, 0], thickness=2):
    right_y_set = []
    right_x_set = []
    right_slope_set = []


    left_y_set = []
    left_x_set = []
    left_slope_set = []


    slope_min = 0.2  # 斜率低阈值
    slope_max = 0.9  # 斜率高阈值，筛选出可能是车道线的直线
    middle_x = image.shape[1] / 2  # 图像中线x坐标
    max_y = image.shape[0]  # 最大y坐标
    # 初始化用于存储右侧车道线和左侧车道线相关数据的列表，
    # 同时定义斜率的上下阈值、图像水平方向的中线坐标以及图像的最大y坐标（高度方向）


    for line in lines:
        for x1, y1, x2, y2 in line:
            fit = np.polyfit((x1, x2), (y1, y2), 1)  # 拟合成直线
            slope = fit[0]  # 斜率
            # 遍历检测到的每一条线段（由坐标(x1, y1)和(x2, y2)表示），
            # 使用numpy的polyfit函数进行一次多项式拟合（即直线拟合），得到直线的斜率


            if slope_min < np.absolute(slope) <= slope_max:
                # 判断斜率的绝对值是否在设定的阈值范围内，筛选出符合要求的直线（可能是车道线）


                # 将斜率大于0且线段X坐标在图像中线右边的点存为右边车道线
                if slope > 0 and x1 > middle_x and x2 > middle_x:
                    right_y_set.append(y1)
                    right_y_set.append(y2)
                    right_x_set.append(x1)
                    right_x_set.append(x2)
                    right_slope_set.append(slope)
                    # 如果直线斜率大于0且其x坐标都在图像中线右侧，将该直线的相关坐标和斜率信息分别添加到右侧车道线对应的列表中


                # 将斜率小于0且线段X坐标在图像中线左边的点存为左边车道线
                elif slope < 0 and x1 < middle_x and x2 < middle_x:
                    left_y_set.append(y1)
                    left_y_set.append(y2)
                    left_x_set.append(x1)
                    left_x_set.append(x2)
                    left_slope_set.append(slope)
                    # 如果直线斜率小于0且其x坐标都在图像中线左侧，将该直线的相关坐标和斜率信息分别添加到左侧车道线对应的列表中


    # 绘制左车道线
    if left_y_set:
        lindex = left_y_set.index(min(left_y_set))  # 最高点
        left_x_top = left_x_set[lindex]
        left_y_top = left_y_set[lindex]
        lslope = np.median(left_slope_set)  # 计算平均值
        # 如果左侧车道线相关的坐标列表不为空，找到y坐标最小（即最高点）的索引，获取对应的x、y坐标作为车道线顶部坐标，
        # 并计算左侧车道线斜率的中位数作为代表斜率


        # 根据斜率计算车道线与图片下方交点作为起点
        left_x_bottom = int(left_x_top + (max_y - left_y_top) / lslope)
        # 根据直线斜率和顶部坐标，计算车道线与图像底部（最大y坐标处）的交点x坐标


        # 绘制线段
        cv.line(image, (left_x_bottom, max_y), (left_x_top, left_y_top), color, thickness)
        # 使用OpenCV的line函数，根据计算得到的起点和顶部坐标，以指定颜色和线宽绘制左侧车道线
    else:
        print("未找到左车道线")
        # 如果左侧车道线坐标列表为空，说明没有检测到左侧车道线，打印提示信息


    # 绘制右车道线
    if right_y_set:
        rindex = right_y_set.index(min(right_y_set))  # 最高点
        right_x_top = right_x_set[rindex]
        right_y_top = right_y_set[rindex]
        rslope = np.median(right_slope_set)
        # 如果右侧车道线相关的坐标列表不为空，找到y坐标最小（即最高点）的索引，获取对应的x、y坐标作为车道线顶部坐标，
        # 并计算右侧车道线斜率的中位数作为代表斜率


        # 根据斜率计算车道线与图片下方交点作为起点
        right_x_bottom = int(right_x_top + (max_y - right_y_top) / rslope)
        # 根据直线斜率和顶部坐标，计算车道线与图像底部（最大y坐标处）的交点x坐标


        # 绘制线段
        cv.line(image, (right_x_top, right_y_top), (right_x_bottom, max_y), color, thickness)
        # 使用OpenCV的line函数，根据计算得到的起点和顶部坐标，以指定颜色和线宽绘制右侧车道线
    else:
        print("未找到右车道线")
        # 如果右侧车道线坐标列表为空，说明没有检测到右侧车道线，打印提示信息


def process_image(image):
    rho = 2  # 霍夫像素单位
    theta = np.pi / 180  # 霍夫角度移动步长
    hof_threshold = 30  # 霍夫平面累加阈值threshold
    min_line_len = 40  # 线段最小长度
    max_line_gap = 80  # 最大允许断裂长度
    # 定义霍夫变换相关的参数，包括像素单位、角度移动步长、累加阈值、线段最小长度以及最大允许断裂长度


    kernel_size = 5  # 高斯滤波器大小size
    canny_low_threshold = 50  # canny边缘检测低阈值
    canny_high_threshold = canny_low_threshold * 3  # canny边缘检测高阈值
    # 定义高斯滤波的卷积核大小以及Canny边缘检测的低阈值和高阈值（高阈值是低阈值的3倍）


    alpha = 0.8  # 原图像权重
    beta = 1.  # 车道线图像权重
    lambda_ = 0.
    # 定义图像融合时原图像和车道线图像的权重系数以及偏置项


    imshape = image.shape  # 获取图像大小
    # 获取输入图像的形状信息（高度、宽度以及通道数等）


    # 灰度图转换
    gray = grayscale(image)
    # 调用grayscale函数将输入的彩色图像转换为灰度图像


    # 高斯滤波
    blur_gray = gaussian_blur(gray, kernel_size)
    # 调用gaussian_blur函数对灰度图像进行高斯滤波操作


    # Canny边缘检测
    edge_image = canny(blur_gray, canny_low_threshold, canny_high_threshold)
    # 调用canny函数对滤波后的灰度图像进行Canny边缘检测


    # 生成Mask掩模
    vertices = np.array([[(0, imshape[0]), (8 * imshape[1] / 20, 11 * imshape[0] / 18),
                       (12 * imshape[1] / 20, 11 * imshape[0] / 18), (imshape[1], imshape[0])]], dtype=np.int32)
    masked_edges = region_of_interest(edge_image, vertices)
    # 定义感兴趣区域的顶点坐标，调用region_of_interest函数生成掩模，得到只保留感兴趣区域的边缘图像


    # 基于霍夫变换的直线检测
    lines = hough_lines(masked_edges, rho, theta, hof_threshold, min_line_len, max_line_gap)
    line_image = np.zeros_like(image)
    # 调用hough_lines函数在感兴趣区域的边缘图像上进行霍夫变换直线检测，
    # 并创建一个与输入图像大小相同、元素全为0的图像，用于后续绘制检测到的直线


    # 绘制车道线线段
    draw_lines(line_image, lines, thickness=10)
    # 调用draw_lines函数在创建的空白图像上根据检测到的直线信息绘制车道线线段


    # 图像融合
    lines_edges = weighted_img(image, line_image, alpha, beta, lambda_)
    return lines_edges
    # 调用weighted_img函数将原图像和绘制了车道线的图像按照设定的权重进行融合，返回融合后的图像


if __name__ == '__main__':
    image = cv.imread('test2b.jpg')
    # 使用OpenCV的imread函数读取指径（'test1.jpg'）的图像文件，将图像加载到内存中
    line_image = process_image(image)
    # 调用process_image函数对读取的图像进行一系列处理，包括灰度转换、滤波、边缘检测、直线检测以及图像融合等操作，得到最终结果图像
    cv.imshow('image', line_image)
    # 使用OpenCV的imshow函数创建一个窗口，并显示处理后的图像，窗口标题为'image'
    cv.waitKey(0)
    # 使用OpenCV的waitKey函数等待用户按键操作，参数0表示一直等待，直到用户按下任意键
